// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/exec/simple_aggregate_adapter.h>
#include <pollux/type/decimal_util.h>

namespace kumo::pollux::functions::aggregate::sparksql {
    /// @tparam TInputType The raw input data type.
    /// @tparam TSumType The type of sum in the output of partial aggregation or the
    /// final output type of final aggregation.
    /// @tparam ResultPrecision The precision of the result type, used for checking
    /// overflow.
    template<typename TInputType, typename TSumType, uint8_t ResultPrecision>
    class DecimalSumAggregate {
    public:
        using InputType = Row<TInputType>;

        using IntermediateType =
        Row</*sum*/ TSumType,
            /*isEmpty*/ bool>;

        using OutputType = TSumType;

        /// Spark's decimal sum doesn't have the concept of a null group, each group
        /// is initialized with an initial value, where sum = 0 and isEmpty = true.
        /// The final agg may fallback to being executed in Spark, so the meaning of
        /// the intermediate data should be consistent with Spark. Therefore, we need
        /// to use the parameter nonNullGroup in writeIntermediateResult to output a
        /// null group as sum = 0, isEmpty = true. nonNullGroup is only available when
        /// default-null behavior is disabled.
        static constexpr bool default_null_behavior_ = false;

        static bool toIntermediate(
            exec::out_type<Row<TSumType, bool> > &out,
            exec::optional_arg_type<TInputType> in) {
            if (in.has_value()) {
                out.copy_from(std::make_tuple(static_cast<TSumType>(in.value()), false));
            } else {
                out.copy_from(std::make_tuple(static_cast<TSumType>(0), true));
            }
            return true;
        }

        /// This struct stores the sum of input values, overflow during accumulation,
        /// and a bool value isEmpty used to indicate whether all inputs are null. The
        /// initial value of sum is 0. We need to keep sum unchanged if the input is
        /// null, as sum function ignores null input. If the isEmpty is true, then it
        /// means there were no values to begin with or all the values were null, so
        /// the result will be null. If the isEmpty is false, then if sum is nullopt
        /// that means an overflow has happened, it returns null.
        struct AccumulatorType {
            std::optional<int128_t> sum{0};
            int64_t overflow{0};
            bool isEmpty{true};

            static constexpr bool is_aligned_ = true;

            AccumulatorType() = delete;

            explicit AccumulatorType(
                HashStringAllocator * /*allocator*/,
                DecimalSumAggregate<TInputType, TSumType, ResultPrecision> * /*fn*/) {
            }

            std::optional<int128_t> computeFinalResult() const {
                if (!sum.has_value()) {
                    return std::nullopt;
                }
                auto const adjustedSum =
                        DecimalUtil::adjustSumForOverflow(sum.value(), overflow);
                if (adjustedSum.has_value() &&
                    DecimalUtil::valueInPrecisionRange(adjustedSum, ResultPrecision)) {
                    return adjustedSum;
                } else {
                    // Found overflow during computing adjusted sum.
                    return std::nullopt;
                }
            }

            bool addInput(
                HashStringAllocator * /*allocator*/,
                exec::optional_arg_type<TInputType> data) {
                if (!data.has_value()) {
                    return false;
                }
                if (!sum.has_value()) {
                    // sum is initialized to 0. When it is nullopt, it implies that the
                    // input data must not be empty.
                    POLLUX_CHECK(!isEmpty);
                    return true;
                }
                int128_t result;
                overflow +=
                        DecimalUtil::addWithOverflow(result, data.value(), sum.value());
                sum = result;
                isEmpty = false;
                return true;
            }

            bool combine(
                HashStringAllocator * /*allocator*/,
                exec::optional_arg_type<Row<TSumType, bool> > other) {
                if (!other.has_value()) {
                    return false;
                }
                auto const otherSum = other.value().template at<0>();
                auto const otherIsEmpty = other.value().template at<1>();

                // isEmpty is never null.
                POLLUX_CHECK(otherIsEmpty.has_value());
                if (isEmpty && otherIsEmpty.value()) {
                    // Both accumulators are empty, no need to do the combination.
                    return false;
                }

                bool currentOverflow = !isEmpty && !sum.has_value();
                bool otherOverflow = !otherIsEmpty.value() && !otherSum.has_value();
                if (currentOverflow || otherOverflow) {
                    sum = std::nullopt;
                    isEmpty = false;
                } else {
                    int128_t result;
                    overflow +=
                            DecimalUtil::addWithOverflow(result, otherSum.value(), sum.value());
                    sum = result;
                    isEmpty &= otherIsEmpty.value();
                }
                return true;
            }

            bool writeIntermediateResult(
                bool nonNullGroup,
                exec::out_type<IntermediateType> &out) {
                if (!nonNullGroup) {
                    // If a group is null, all values in this group are null. In Spark, this
                    // group will be the initial value, where sum is 0 and isEmpty is true.
                    out = std::make_tuple(static_cast<TSumType>(0), true);
                } else {
                    auto finalResult = computeFinalResult();
                    if (finalResult.has_value()) {
                        out = std::make_tuple(
                            static_cast<TSumType>(finalResult.value()), isEmpty);
                    } else {
                        // Sum should be set to null on overflow,
                        // and isEmpty should be set to false.
                        out.template set_null_at<0>();
                        out.template get_writer_at<1>() = false;
                    }
                }
                return true;
            }

            bool writeFinalResult(bool nonNullGroup, exec::out_type<OutputType> &out) {
                if (!nonNullGroup || isEmpty) {
                    // If isEmpty is true, we should set null.
                    return false;
                }
                auto finalResult = computeFinalResult();
                if (finalResult.has_value()) {
                    out = static_cast<TSumType>(finalResult.value());
                    return true;
                } else {
                    // Sum should be set to null on overflow.
                    return false;
                }
            }
        };
    };
} // namespace kumo::pollux::functions::aggregate::sparksql
